import asyncio
from langchain_mcp_adapters.client import MultiServerMCPClient, load_mcp_tools
from contextlib import AsyncExitStack

class MCPClientWrapper:
    def __init__(self):
        self.task = None
        self.tools = []

    async def cleanup(self):
        if self.task:
            self.task.cancel()
            await asyncio.gather(self.task, return_exceptions=True)
        self.task = None
    
    async def connect(self, mcp_configs) -> str:
        await self.cleanup()
        in_queue = asyncio.Queue()
        out_queue = asyncio.Queue()
        self.task = asyncio.create_task(self._connect(in_queue, out_queue))
        await in_queue.put(mcp_configs)
        await in_queue.join()
        self.tools = await out_queue.get()
        out_queue.task_done()
        tool_names = [t.name for t in self.tools]
        return f"Connected to MCP. Available tools: {tool_names}"

    async def _connect(self, in_channel:asyncio.Queue, out_channel:asyncio.Queue) -> str:
        mcp_config = await in_channel.get()
        in_channel.task_done()
        client = MultiServerMCPClient(mcp_config)
        tools = []
        
        async with AsyncExitStack() as stack:
            for k in client.connections.keys():
                session = await stack.enter_async_context(client.session(k))
                tools += await load_mcp_tools(session)
            await out_channel.put(tools)
            await out_channel.join()
            try:
                while True:
                    await asyncio.sleep(10)
            except asyncio.CancelledError:
                print("cancelled the old server")
            finally:
                print("stop the old server")